
import os.path
import time
import subprocess
import datetime
import re
import shutil
from utils import ssh_scp_files
import scp

dct = {'Aquila-7b-20n8g-V2': "./Aquila-7b-20n8g.yaml.log.txt",
        'Aquila-7b-24n8g': "./rank0.log",
        '30B': "/share/project/64node-bmt-flashatten/log",
        'Aquila-7b-8n8g-code':'Aquila-7b-8n8g-code.yaml.log.txt',
        'Aquila-7b-16n8g': 'Aquila-7b-16n8g.yaml.log.txt.9.192.168.25.2'}
prev_loss = 0
loss = -1

# src_file = "/share/project/64node-ds/current_training/tboard4000/llama_node64_zero2_mp4_bs4_warm0.2_alldatav1_0413/events.out.tfevents.1681688514.job-cf2f03fc-2e3f-4b8f-82f7-0526efbe98c0-master-0.233966.0"
# dst_file = "/share/project/64node-ds/current_training/tboard4000/test/events.out.tfevents.1681688514.job-cf2f03fc-2e3f-4b8f-82f7-0526efbe98c0-master-0.233966.0"

checked_models = ['30B','Aquila-7b-24n8g']
checked_models = ['30B','Aquila-7b-8n8g-code']
prev_iter_7B = '-1'
prev_iter_7B_21n8g = '-1'
prev_iter_30B = '-1'
prev_iter_7b_8n8g_code = '-1'



def send_message(model_name, wandb_link, ssh_host, ssh_user, ssh_password, ssh_port, local_volume, remote_volume):
        model_name = 'Aquila-7b-8n8g-code'
        connected = True 
        try:
            ssh_scp_files(ssh_host, ssh_user, ssh_password, ssh_port, local_volume, remote_volume)
        except scp.SCPException:
            connected = False
        if os.path.exists(dct[model_name]) and connected:
            with open(dct[model_name], 'r') as f:
                last_line = None
                lines = [line.rstrip('\n') for line in f.readlines()]
            # Get the modification time of the file and current time
            

            file_modified_time = os.path.getmtime(dct[model_name])
            current_time = time.time()
            
            # Calculate the time difference in seconds
            time_diff = current_time - file_modified_time

            # script for 7B
            time_print = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            print('time',time_print)
            latest = True
            for i in reversed(range(len(lines))):
                try:
                    tmp = [e.strip() for e in lines[i].split('|')]
                    loss_str = [e for e in tmp if e.startswith('loss')][0]
                except IndexError:
                    continue

                pattern = r'iteration\s(.*?)(?=\|)'
                match = re.search(pattern, lines[i])
                if match:
                    content = match.group(0)

                loss = float(loss_str.split()[1])
                break # Once find a info line, exit

            cur_iter = content.split()[1][:-1]
            if loss != -1:
                if cur_iter == prev_iter_7b_8n8g_code:
                    print("The file was not modified within the last 60 minutes.")
                    # Run the curl command to send a warning
                    curl_command = """curl --location --request POST 'https://flagopen.baai.ac.cn/warning/send_warning' \
                            --header 'Content-Type: application/json' \
                            --data-raw '{
                                "text": "警告：Aquila-7b-8n8g-code模型iteration无变化，怀疑训练中止; 当前时间为""" + str(time_print) + """,告警信息为"""+lines[-1]+""" ",
                                "user_list": ["g8b45d2g"]
                            }'"""
                    subprocess.run(curl_command, shell=True)
                    checked_models.remove(model_name)
                elif loss == 0:
                    curl_command = """curl --location --request POST 'https://flagopen.baai.ac.cn/warning/send_warning' \
                            --header 'Content-Type: application/json' \
                            --data-raw '{
                                "text": "警告：Aquila-7b-8n8g-code模型loss为0，怀疑训练出现问题; 当前时间为""" + str(time_print) + """ ",
                                "user_list": ["gb8bgd55"]
                            }'"""
                    subprocess.run(curl_command, shell=True)
                    checked_models.remove(model_name) 
                else:
                    curl_command = """curl --location --request POST 'https://flagopen.baai.ac.cn/warning/send_message' \
                            --header 'Content-Type: application/json' \
                            --data-raw '{
                                "text": "Aquila-7b-8n8g-code模型训练正在进行中; 当前时间为""" + str(time_print) + """, \\n               进程： """+content+""", \\n               当前loss为 """+str(loss)+""" \\n               https://wandb.ai/ldwang/Aquila-7b-8n8g-code/runs/uelm2qug?workspace=user-",
                                "user_list": ["9gdf9d4e"]
                            }'"""
                prev_loss = loss
                prev_iter_7b_8n8g_code = cur_iter
            else:
                curl_command = """curl --location --request POST 'https://flagopen.baai.ac.cn/warning/send_message' \
                        --header 'Content-Type: application/json' \
                        --data-raw '{
                            "text": "Aquila-7b-8n8g-code模型训练正在进行中; 当前时间为""" + str(time_print) + """, \\n               进程： 尚未有任何loss信息",
                            "user_list": ["9gdf9d4e"]
                        }'"""    
            subprocess.run(curl_command, shell=True)
        else:
            curl_command = """curl --location --request POST 'https://flagopen.baai.ac.cn/warning/send_message' \
                    --header 'Content-Type: application/json' \
                    --data-raw '{
                        "text": "未能连接模型训练的远程服务器; 当前时间为""" + str(time_print) + """",
                    }'"""   
            subprocess.run(curl_command, shell=True)




while True:
    # shutil.copy(src_file, dst_file)  # for tensorboard
    if 'Aquila-7b-8n8g-code' in checked_models:
        model_name = 'Aquila-7b-8n8g-code'
        connected = True 
        try:
            # ssh_scp_files('36.103.236.181', 'ubuntu', 'p5dkf@v2QvSGm@bW%aTc', 22, './', '/data/init_training_old/bmtrain/logs/rank0.log')
            ssh_scp_files('36.103.236.163', 'ubuntu', 'dzZgKaR#xTeL6Dgw2@HtD', 22, './', '/data2/ldwang/checkpoints/Aquila-7b-8n8g-code/2023042814/Aquila-7b-8n8g-code.yaml.log.txt')
        except scp.SCPException:
            connected = False
        if os.path.exists(dct[model_name]) and connected:
            with open(dct[model_name], 'r') as f:
                last_line = None
                lines = [line.rstrip('\n') for line in f.readlines()]
            # Get the modification time of the file and current time
            

            file_modified_time = os.path.getmtime(dct[model_name])
            current_time = time.time()
            
            # Calculate the time difference in seconds
            time_diff = current_time - file_modified_time
            import pdb;pdb.set_trace()
            # script for 7B
            time_print = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            print('time',time_print)
            latest = True
            for i in reversed(range(len(lines))):
                try:
                    tmp = [e.strip() for e in lines[i].split('|')]
                    loss_str = [e for e in tmp if e.startswith('loss')][0]
                except IndexError:
                    continue

                pattern = r'iteration\s(.*?)(?=\|)'
                match = re.search(pattern, lines[i])
                if match:
                    content = match.group(0)

                loss = float(loss_str.split()[1])
                break # Once find a info line, exit

            cur_iter = content.split()[1][:-1]
            if loss != -1:
                if cur_iter == prev_iter_7b_8n8g_code:
                    print("The file was not modified within the last 60 minutes.")
                    # Run the curl command to send a warning
                    curl_command = """curl --location --request POST 'https://flagopen.baai.ac.cn/warning/send_warning' \
                            --header 'Content-Type: application/json' \
                            --data-raw '{
                                "text": "警告：Aquila-7b-8n8g-code模型iteration无变化，怀疑训练中止; 当前时间为""" + str(time_print) + """,告警信息为"""+lines[-1]+""" ",
                                "user_list": ["g8b45d2g"]
                            }'"""
                    subprocess.run(curl_command, shell=True)
                    checked_models.remove(model_name)
                elif loss == 0:
                    curl_command = """curl --location --request POST 'https://flagopen.baai.ac.cn/warning/send_warning' \
                            --header 'Content-Type: application/json' \
                            --data-raw '{
                                "text": "警告：Aquila-7b-8n8g-code模型loss为0，怀疑训练出现问题; 当前时间为""" + str(time_print) + """ ",
                                "user_list": ["gb8bgd55"]
                            }'"""
                    subprocess.run(curl_command, shell=True)
                    checked_models.remove(model_name) 
                else:
                    curl_command = """curl --location --request POST 'https://flagopen.baai.ac.cn/warning/send_message' \
                            --header 'Content-Type: application/json' \
                            --data-raw '{
                                "text": "Aquila-7b-8n8g-code模型训练正在进行中; 当前时间为""" + str(time_print) + """, \\n               进程： """+content+""", \\n               当前loss为 """+str(loss)+""" \\n               https://wandb.ai/ldwang/Aquila-7b-8n8g-code/runs/uelm2qug?workspace=user-",
                                "user_list": ["9gdf9d4e"]
                            }'"""
                prev_loss = loss
                prev_iter_7b_8n8g_code = cur_iter
            else:
                curl_command = """curl --location --request POST 'https://flagopen.baai.ac.cn/warning/send_message' \
                        --header 'Content-Type: application/json' \
                        --data-raw '{
                            "text": "Aquila-7b-8n8g-code模型训练正在进行中; 当前时间为""" + str(time_print) + """, \\n               进程： 尚未有任何loss信息",
                            "user_list": ["9gdf9d4e"]
                        }'"""    
            subprocess.run(curl_command, shell=True)
        else:
            print("The file path does not exist.")

    if 'Aquila-7b-24n8g' in checked_models:
        connected = True 
        try:
            # ssh_scp_files('36.103.236.181', 'ubuntu', 'p5dkf@v2QvSGm@bW%aTc', 22, './', '/data/init_training_old/bmtrain/logs/rank0.log')
            ssh_scp_files('36.103.236.162', 'ubuntu', 'Dg%xoL%tDaCJSg4I', 22, './', '/data2/ldwang/checkpoints/Aquila-7b-24n8g/2023042722/rank0.log')
        except scp.SCPException:
            connected = False
        if os.path.exists(dct['Aquila-7b-24n8g']) and connected:
            with open(dct['Aquila-7b-24n8g'], 'r') as f:
                last_line = None
                lines = [line.rstrip('\n') for line in f.readlines()]
            # Get the modification time of the file and current time
            

            file_modified_time = os.path.getmtime(dct['Aquila-7b-24n8g'])
            current_time = time.time()
            
            # Calculate the time difference in seconds
            time_diff = current_time - file_modified_time

            # script for 7B
            time_print = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            print('time',time_print)
            latest = True
            for i in reversed(range(len(lines))):
                try:
                    tmp = [e.strip() for e in lines[i].split('|')]
                    loss_str = [e for e in tmp if e.startswith('loss')][0]
                except IndexError:
                    continue

                pattern = r'iteration\s(.*?)(?=\|)'
                match = re.search(pattern, lines[i])
                if match:
                    content = match.group(0)

                loss = float(loss_str.split()[1])
                break # Once find a info line, exit

            cur_iter_7B_21n8g = content.split()[1][:-1]
            if loss != -1:
                if cur_iter_7B_21n8g == prev_iter_7B_21n8g:
                    print("The file was not modified within the last 60 minutes.")
                    # Run the curl command to send a warning
                    curl_command = """curl --location --request POST 'https://flagopen.baai.ac.cn/warning/send_warning' \
                            --header 'Content-Type: application/json' \
                            --data-raw '{
                                "text": "警告：Aquila-7b-24n8g模型iteration无变化，怀疑训练中止; 当前时间为""" + str(time_print) + """,告警信息为"""+lines[-1]+""" ",
                                "user_list": ["gb8bgd55"]
                            }'"""
                    subprocess.run(curl_command, shell=True)
                    checked_models.remove('Aquila-7b-24n8g')
                elif loss == 0:
                    curl_command = """curl --location --request POST 'https://flagopen.baai.ac.cn/warning/send_warning' \
                            --header 'Content-Type: application/json' \
                            --data-raw '{
                                "text": "警告：Aquila-7b-24n8g模型loss为0，怀疑训练出现问题; 当前时间为""" + str(time_print) + """ ",
                                "user_list": ["gb8bgd55"]
                            }'"""
                    subprocess.run(curl_command, shell=True)
                    checked_models.remove('Aquila-7b-24n8g') 
                else:
                    curl_command = """curl --location --request POST 'https://flagopen.baai.ac.cn/warning/send_message' \
                            --header 'Content-Type: application/json' \
                            --data-raw '{
                                "text": "Aquila-7b-24n8g模型训练正在进行中; 当前时间为""" + str(time_print) + """, \\n               进程： """+content+""", \\n               当前loss为 """+str(loss)+""" \\n               lr&loss曲线链接: https://wandb.ai/ldwang/Aquila-7b-24n8g/runs/126312ck",
                                "user_list": ["9gdf9d4e"]
                            }'"""
                prev_loss = loss
                prev_iter_7B_21n8g = cur_iter_7B_21n8g
            else:
                curl_command = """curl --location --request POST 'https://flagopen.baai.ac.cn/warning/send_message' \
                        --header 'Content-Type: application/json' \
                        --data-raw '{
                            "text": "Aquila-7b-24n8g模型训练正在进行中; 当前时间为""" + str(time_print) + """, \\n               进程： 尚未有任何loss信息",
                            "user_list": ["9gdf9d4e"]
                        }'"""    
            subprocess.run(curl_command, shell=True)
        else:
            print("The file path does not exist.")

    if 'Aquila-7b-20n8g-V2' in checked_models:
        connected = True 
        try:
            ssh_scp_files(ssh_host='36.103.236.162', ssh_user='ubuntu', ssh_password='Dg%xoL%tDaCJSg4I', ssh_port=22, local_volume='./', remote_volume='/data2/ldwang/checkpoints/Aquila-7b-20n8g/2023042608/Aquila-7b-20n8g.yaml.log.txt')
        except scp.SCPException:
            connected = False

        if os.path.exists(dct['Aquila-7b-20n8g-V2']) and connected:
            with open(dct['Aquila-7b-20n8g-V2'], 'r') as f:
                last_line = None
                lines = [line.rstrip('\n') for line in f.readlines()]
            # Get the modification time of the file and current time
            

            file_modified_time = os.path.getmtime(dct['Aquila-7b-20n8g-V2'])
            current_time = time.time()
            
            # Calculate the time difference in seconds
            time_diff = current_time - file_modified_time

            # script for 7B
            time_print = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            print('time',time_print)
            latest = True
            for i in reversed(range(len(lines))):
                try:
                    tmp = [e.strip() for e in lines[i].split('|')]
                    loss_str = [e for e in tmp if e.startswith('loss')][0]
                except IndexError:
                    continue

                pattern = r'iteration\s(.*?)(?=\|)'
                match = re.search(pattern, lines[i])
                if match:
                    content = match.group(0)

                try:
                    loss = float(loss_str.split()[1])
                except ValueError:
                    continue
                break # Once find a info line, exit

            cur_iter_7B = content.split()[1][:-1]
            if loss != -1:
                if cur_iter_7B == prev_iter_7B:
                    print("The file was not modified within the last 60 minutes.")
                    # Run the curl command to send a warning
                    curl_command = """curl --location --request POST 'https://flagopen.baai.ac.cn/warning/send_warning' \
                            --header 'Content-Type: application/json' \
                            --data-raw '{
                                "text": "警告：Aquila-7b-20n8g-V2模型iteration无变化，怀疑训练中止; 当前时间为""" + str(time_print) + """,告警信息为"""+lines[-1]+""" ",
                                "user_list": ["2dc93b7c"]
                            }'"""
                    subprocess.run(curl_command, shell=True)
                    checked_models.remove('Aquila-7b-20n8g-V2')
                elif loss == 0:
                    curl_command = """curl --location --request POST 'https://flagopen.baai.ac.cn/warning/send_warning' \
                            --header 'Content-Type: application/json' \
                            --data-raw '{
                                "text": "警告：Aquila-7b-20n8g-V2模型loss为0，怀疑训练出现问题; 当前时间为""" + str(time_print) + """ ",
                                "user_list": ["2dc93b7c"]
                            }'"""
                    subprocess.run(curl_command, shell=True)
                    checked_models.remove('Aquila-7b-20n8g-V2') 
                else:
                    curl_command = """curl --location --request POST 'https://flagopen.baai.ac.cn/warning/send_message' \
                            --header 'Content-Type: application/json' \
                            --data-raw '{
                                "text": "Aquila-7b-20n8g-V2模型训练正在进行中; 当前时间为""" + str(time_print) + """, \\n               进程： """+content+""", \\n               当前loss为 """+str(loss)+""" ",
                                "user_list": ["9gdf9d4e"]
                            }'"""
                prev_loss = loss
                prev_iter_7B = cur_iter_7B
            else:
                curl_command = """curl --location --request POST 'https://flagopen.baai.ac.cn/warning/send_message' \
                        --header 'Content-Type: application/json' \
                        --data-raw '{
                            "text": "Aquila-7b-20n8g-V2模型训练正在进行中; 当前时间为""" + str(time_print) + """, \\n               进程： 尚未有任何loss信息",
                            "user_list": ["9gdf9d4e"]
                        }'"""    
            subprocess.run(curl_command, shell=True)
        else:
            print("The file path does not exist.")
    if '30B' in checked_models:
        if os.path.exists(dct['30B']):
            current_time = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            with open(dct['30B'], 'r') as f:
                last_line = None
                lines = [line.rstrip('\n') for line in f.readlines()]
            # Get the modification time of the file and current time

            file_modified_time = os.path.getmtime(dct['30B'])
            current_time = time.time()
            
            # Calculate the time difference in seconds
            time_diff = current_time - file_modified_time
            print(time_diff)
            time_print = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            # Check if the file was modified within the last 60 minutes (3600 seconds)
            
            if time_diff <= 3600:
                print("The file was modified within the last 60 minutes.")
                losses = []
                count = 0
                latest = True
                content = ''
                for i in reversed(range(len(lines))):
                    try:
                        tmp = [e.strip() for e in lines[i].split('|')]
                        loss_str = [e for e in tmp if e.startswith('loss')][0]
                    except IndexError:
                        continue
                    if latest:
                        pattern = r'iteration\s(.*?)(?=\|)'
                        match = re.search(pattern, lines[i])
                        if match:
                            content = match.group(0)
                        latest = False
                    try:
                        loss = float(loss_str.split()[1])
                    except ValueError:
                        continue
                    losses.append(loss)
                    count += 1
                    if count == 50:
                        break
                if content:
                    cur_iter_30B = content.split()[1][:-1]
                    if cur_iter_30B == prev_iter_30B:
                        print("The file was not modified within the last 60 minutes.")
                        # Run the curl command to send a warning
                        curl_command = """curl --location --request POST 'https://flagopen.baai.ac.cn/warning/send_warning' \
                                --header 'Content-Type: application/json' \
                                --data-raw '{
                                    "text": "警告：Aquila30B模型iteration无变化，怀疑训练中止; 当前时间为""" + str(time_print) + """,告警信息为"""+lines[-1]+""" ",
                                    "user_list": ["2dc93b7c"]
                                }'"""
                        subprocess.run(curl_command, shell=True)
                        checked_models.remove('30B')
                    elif loss == 0:
                        curl_command = """curl --location --request POST 'https://flagopen.baai.ac.cn/warning/send_warning' \
                                --header 'Content-Type: application/json' \
                                --data-raw '{
                                    "text": "警告：Aquila30B模型loss为0，怀疑训练出现问题; 当前时间为""" + str(time_print) + """ ",
                                    "user_list": ["2dc93b7c"]
                                }'"""
                        subprocess.run(curl_command, shell=True)
                        checked_models.remove('30B')                      
                    else:
                        loss = str(sum(losses)/len(losses))
                        curl_command = """curl --location --request POST 'https://flagopen.baai.ac.cn/warning/send_message' \
                                --header 'Content-Type: application/json' \
                                --data-raw '{
                                    "text": "Aquila30B模型训练正在进行中; 当前时间为""" + str(time_print) + """, \\n               进程： """+content+""", \\n               过去"""+str(len(losses))+"""步平均loss为 """+loss+"""  \\n               lr&loss曲线链接: https://wandb.ai/aaronlg2010/Aquila-30b-64n8g-from-scratch?workspace=user-aaronlg2010",
                                    "user_list": ["9gdf9d4e"]
                                }'"""
                    prev_iter_30B = cur_iter_30B
                    subprocess.run(curl_command, shell=True)
                else:
                    curl_command = """curl --location --request POST 'https://flagopen.baai.ac.cn/warning/send_message' \
                            --header 'Content-Type: application/json' \
                            --data-raw '{
                                "text": "Aquila30B模型训练正在进行中; 当前时间为""" + str(time_print) + """, \\n               进程： 目前暂无loss输出",
                                "user_list": ["9gdf9d4e"]
                            }'"""
                    subprocess.run(curl_command, shell=True)
                    prev_time = current_time


            else:
                print("The file was not modified within the last 60 minutes.")
                # Run the curl command to send a warning
                curl_command = """curl --location --request POST 'https://flagopen.baai.ac.cn/warning/send_warning' \
                        --header 'Content-Type: application/json' \
                        --data-raw '{
                            "text": "警告：Aquila30B模型一小时未更新，怀疑训练中止; 当前时间为""" + str(time_print) + """ ",
                            "user_list": ["2dc93b7c"]
                        }'"""
                # import pdb;pdb.set_trace()
                subprocess.run(curl_command, shell=True)
                checked_models.remove('30B')
                
        else:
            print("The file path does not exist.")
        # script for 30B
    if not checked_models:
        break

    time.sleep(3500)

